#!/bin/env python

import cv2
import numpy as np
import gzip

show_x = 80
show_y = 40

def load_mnist_images(path):
    print(path)
    with gzip.open(path, 'rb') as f:
        # 读取魔数
        magic_number = int.from_bytes(f.read(4), 'big')
        # 读取数据的数量
        num_images = int.from_bytes(f.read(4), 'big')
        # 读取维度信息（对于MNIST图像数据，维度是1x28x28）
        num_rows = int.from_bytes(f.read(4), 'big')
        num_cols = int.from_bytes(f.read(4), 'big')

        print("magic_number:", magic_number)
        print("num_images:", num_images)
        print("num_rows:", num_rows)
        print("num_cols:", num_cols)

        # 读取图像数据
        buf = f.read()
        # 将字节数据转换为numpy数组
        data = np.frombuffer(buf, dtype=np.uint8)
        # 重新整形为图像数组
        data = data.reshape(-1, num_rows, num_cols)
        return data

def load_mnist_labels(path):
    print(path)
    with gzip.open(path, 'rb') as f:
        # 读取魔数
        magic_number = int.from_bytes(f.read(4), 'big')
        # 读取数据的数量
        num_images = int.from_bytes(f.read(4), 'big')

        print("magic_number:", magic_number)
        print("num_images:", num_images)

        # 读取标签数据
        buf = f.read()
        # 将字节数据转换为numpy数组
        labels = np.frombuffer(buf, dtype=np.uint8)
        return labels

# 指定文件路径
train_images_path = '../../resource/data/MNIST/train-images-idx3-ubyte.gz'
train_labels_path = '../../resource/data/MNIST/train-labels-idx1-ubyte.gz'
test_images_path = '../../resource/data/MNIST/t10k-images-idx3-ubyte.gz'
test_labels_path = '../../resource/data/MNIST/t10k-labels-idx1-ubyte.gz'

# 加载训练集图像和标签
train_images = load_mnist_images(train_images_path)
train_labels = load_mnist_labels(train_labels_path)

# 加载测试集图像和标签
test_images = load_mnist_images(test_images_path)
test_labels = load_mnist_labels(test_labels_path)

# 打印数据集信息
print("训练集样本数量:", train_images.shape[0])
print("测试集样本数量:", test_images.shape[0])
print("输入特征形状:", train_images[0].shape)
print("标签形状:", train_labels.shape)

for i in range(0, 10):
    print(train_labels[i])

windowname = "digital"

cv2.namedWindow(windowname, cv2.WINDOW_NORMAL)
cv2.resizeWindow(windowname, train_images[0].shape[1] * show_x, train_images[0].shape[0] * show_y)

for i in range(0, show_y):
    horizontal_stack = train_images[i * 0 + i]
    for j in range(1, show_x):
        horizontal_stack = np.hstack((horizontal_stack, train_images[i * 10 + j]))
    if i == 0:
        vertical_stack = horizontal_stack
    else:
        vertical_stack  = np.vstack((vertical_stack, horizontal_stack))

cv2.imshow(windowname, vertical_stack)

cv2.waitKey(0)

print('销毁窗口')
cv2.destroyAllWindows()
